# -*- coding: utf-8 -*- 
# @Time : 2022/4/3 20:36 
# @Author : zzuxyj 
# @File : 09-nn-rule.py


"""
非线性激活函数

"""


import torchvision
import torch

#数据加载
from torch import nn
from torch.nn import ReLU, Sigmoid ,Softmax
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

testSet = torchvision.datasets.CIFAR10("../dataset/CIFAR10" , download=True , train=False , transform=torchvision.transforms.ToTensor())
dataloader = DataLoader(testSet , batch_size=64)


# 模型
class Modle(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.rule = ReLU()
        self.sigmoid = Sigmoid()

    def forward(self , input):
        output = self.rule(input)
        output = self.sigmoid(output)
        return output

def train():
    # 定义模型
    model = Modle()

    # 遍历数据 , tensorboard
    writer = SummaryWriter("logs09")
    step = 0
    for data in dataloader:
        imgs, target = data
        writer.add_images("RuelBefore", imgs, global_step=step)
        output = model(imgs)
        writer.add_images("RuelAfter", output, global_step=step)
        step += 1

    writer.close()

if __name__ == '__main__':
    test = torch.tensor([
        [1,-0.5],[-1,3]
    ])
    print(test.shape)
    test = torch.reshape(test , (-1,1,2,2))
    print(test.shape)


    """
    torch.Size([2, 2])
    torch.Size([1, 1, 2, 2])
    """